#include <iostream>
#include <glad/glad.h>
#include <sstream>
#include "GLFW/glfw3.h"
#include "../gl/GLShaders.h"
#include "../cpu/WritePng.h"

#define NUM_GRID_LUT_SRC 17
#define IMAGE_BIT_LENGTH 10
#define MSM_MEDIA_ALIGN(__sz, __align) (((__align) & ((__align) - 1)) ?\
	((((__sz) + (__align) - 1) / (__align)) * (__align)) :\
	(((__sz) + (__align) - 1) & (~((__align) - 1))))

uint32_t windowWidth = 800;
uint32_t windowHeight = 600;
float filterRatio = 1.f;
GLuint VAO;

static float CUBE[] = {
    #include "S_Night_Time_2_17.cube"
};

const char *m_pVertexShaderSource =
        "#version 330         \n"
        "layout (location = 0) in vec2 a_position;\n"
        "layout (location = 1) in vec2 a_texcoord;\n"
        "out vec2 v_texCoord;\n"
        "void main()\n"
        "{\n"
        "    gl_Position = vec4(a_position, 0, 1);\n"
        "    v_texCoord = a_texcoord;\n"
        "}\n";

const char *m_pFragmentShaderSource =
        "#version 330\n"
        "in vec2 v_texCoord;\n"
        "uniform usampler2D yTexture;\n"
        "uniform usampler2D uvTexture;\n"
        "uniform sampler3D lutTexture;\n"
        "uniform float filterRatio;\n"
        "uniform uint uvoffset;\n"
        "out vec4 FragColor;\n"
        "\n"
        "uvec3 tex_yuv(ivec2 pos){\n"
        "    uvec3 yuv;\n"
        "    yuv.x = texelFetch(yTexture, pos, 0).r;\n"
        "    yuv.y = texelFetch(uvTexture, pos / 2, 0).g;\n"
        "    yuv.z = texelFetch(uvTexture, pos / 2, 0).r;\n"
        "    return yuv;\n"
        "}\n"
        "\n"
        "vec3 bt_10bit_full_yuv_to_rgb(vec3 yuv){\n"
        "    return mat3(\n"
        "        1.000000, 1.0000000, 1.00000,\n"
        "        0.000000, -0.344000, 1.77200,\n"
        "        1.402000, -0.714000, 0.00000\n"
        "    ) * (yuv - vec3(0.0, uvoffset, uvoffset));\n"
        "}\n"
        "\n"
        "vec3 bt_10bit_full_rgb_to_yuv(vec3 rgb){\n"
        "    vec3 yuv = mat3(\n"
        "        0.299000, -0.169000, 0.50000,\n"
        "        0.587000, -0.331000, -0.4190,\n"
        "        0.114000, 0.5000000, -0.0810\n"
        "    ) * rgb;\n"
        "    yuv.y += uvoffset;\n"
        "    yuv.z += uvoffset;\n"
        "    return yuv;\n"
        "}\n"
        "\n"
        "vec3 lut_cube_rgb(vec3 rgb){\n"
        "    return vec3(mix(rgb, texture(lutTexture, rgb).rgb, filterRatio));\n"
        "}\n"
        "\n"
        "void main() {\n"
        "    ivec2 yTexCoordInt = ivec2(v_texCoord * textureSize(yTexture, 0));\n"
        "    uvec3 pre_yuv16 = tex_yuv(yTexCoordInt);\n"
        "    vec3 pre_yuv10 = pre_yuv16 >> 6;\n"
        "    vec3 rgb10 = bt_10bit_full_yuv_to_rgb(pre_yuv10);\n"
        "    vec3 rgb = rgb10 / 1023;\n"
        "    vec3 lut_rgb = lut_cube_rgb(rgb);\n"
//        "    uvec3 post_yuv10 = uvec3(bt_10bit_full_rgb_to_yuv(lut_rgb) * 1023);\n"
//        "    uvec3 post_yuv16 = post_yuv10 << 6;\n"
//        "    FragColor = uvec4(pre_yuv16, 65535);\n"
        "    FragColor = vec4(lut_rgb, 1.0);\n"
        "}";

void initMesh(GLuint program){
    float x_scale = 1.0;
    float y_scale = 1.0;
    GLfloat vertices[] = {
            -1.0f * x_scale, -1.0f * y_scale,
            1.0f * x_scale, -1.0f * y_scale,
            -1.0f * x_scale, 1.0f * y_scale,
            1.0f * x_scale, 1.0f * y_scale,
    };

    GLfloat texCoords[] = {
            0.0f, 0.0f,
            1.0f, 0.0f,
            0.0f, 1.0f,
            1.0f, 1.0f
    };
    glGenVertexArrays(1, &VAO);
    GLuint VBO[2];
    glGenBuffers(2, VBO);

    glBindVertexArray(VAO);
    glBindBuffer(GL_ARRAY_BUFFER, VBO[0]);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);
    GLuint  loc = glGetAttribLocation(program, "a_position");
    glEnableVertexAttribArray(loc);
    glVertexAttribPointer(loc, 2, GL_FLOAT, false, 0, (void*) 0);

    glBindBuffer(GL_ARRAY_BUFFER, VBO[1]);
    glBufferData(GL_ARRAY_BUFFER, sizeof(texCoords), texCoords, GL_STATIC_DRAW);
    loc = glGetAttribLocation(program, "a_texcoord");
    glEnableVertexAttribArray(loc);
    glVertexAttribPointer(loc, 2, GL_FLOAT, false, 0, (void*) 0);
    glBindVertexArray(0);
}

/* Load YUV data:: Actually, we hope the YUV data to use is output from IPE. */
uint16_t*  loadImage(int* width, int* height)
{
#if 1
    const uint32_t imageWidth = 3680;
    const uint32_t imageHeight = 2760;
    uint32_t stride = MSM_MEDIA_ALIGN(imageWidth * 2, 256);
    uint32_t scanline = MSM_MEDIA_ALIGN(imageHeight, 32);
    uint32_t srcSize = stride * scanline * 3 / 2;

    FILE *fp = fopen("cameraInput_3680x2760_2.p010", "rb");
    if (!fp) {
        printf("oepn file fail \n");
        return NULL;
    }
    uint16_t* srcBuf = (uint16_t*)malloc(srcSize);
    fread(srcBuf, 1, srcSize, fp);
    fclose(fp);
#else
    /***** Test image data -> ******/
    /*  */
    const int imageWidth = 2;
    const int imageHeight = 3;
    /* YUV Image  ##Because this is a test data, it has only 2x2 pixels.*/
    const uint16_t yuvTestData[imageHeight*imageWidth][3] = {
            {117,552,479}, {297,500, 530},
            {405,371,788}, {529,596, 398},
            {878,101,443}, {911,267, 103}
    };
    /***** <- Test image data ******/

    int index = 0;
    uint16_t* yuvSrc = new uint16_t[imageWidth * imageHeight * 3];
    for (int y = 0; y < imageHeight; y++) {
        for (int x = 0; x < imageWidth; x++) {
            for (int i = 0; i < 3; i++) yuvSrc[index++] = yuvTestData[y * imageWidth + x][i];
        }
    }
#endif
    *width = stride / 2;
    *height = scanline;
    return srcBuf;
}

bool writeRgb(const char *filePath, uint32_t width, uint32_t height, size_t size, uint8_t *rgbBuffer, bool isGenPng){
    FILE *fp = fopen(filePath,"wb");
    if (fp){
        if(isGenPng) {
            std::stringstream ss;
            ss << filePath << ".png";
            FILE *pngFp = fopen(ss.str().c_str(),"wb");
            if(pngFp){
                svpng(pngFp, width, height, rgbBuffer, 0);
                fclose(pngFp);
                pngFp = nullptr;
            }
        }
        fwrite(rgbBuffer, size, 1, fp);
        fclose(fp);
        fp = nullptr;
        return true;
    } else {
        std::cerr << "Could not read rgb file:" << filePath << std::endl;
        return false;
    }
}

int main() {
    std::cout << "Lut yuv" << std::endl;
    if (!glfwInit()) {
        std::cerr << "Failed init glfw!" << std::endl;
        return -1;
    }
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    GLFWwindow *pWindow = glfwCreateWindow(windowWidth, windowHeight, "Lut yuv", nullptr, nullptr);
    if (!pWindow) {
        std::cerr << "Failed create glfw window!" << std::endl;
        return -1;
    }
    glfwMakeContextCurrent(pWindow);
    glfwSwapInterval(0);
    glfwSetFramebufferSizeCallback(pWindow, [](GLFWwindow *window, int w, int h) {
        windowWidth = w;
        windowHeight = h;
    });
    glfwSetKeyCallback(pWindow, [](GLFWwindow* window, int key, int scancode, int action, int mods){
        if(key == GLFW_KEY_UP && (action == GLFW_PRESS || action == GLFW_REPEAT)){
            if(filterRatio < 1.f){
                filterRatio += 0.1f;
            }
        }
        if(key == GLFW_KEY_DOWN && (action == GLFW_PRESS || action == GLFW_REPEAT)){
            if(filterRatio > 0.f){
                filterRatio -= 0.1f;
            }
        }
    });

    // 3.glad管理函数指针
    bool bLoadSuc = gladLoadGLLoader((GLADloadproc) glfwGetProcAddress);
    if (!bLoadSuc) {
        std::cout << "Failed load glProcAddress!" << std::endl;
        return -1;
    }

    GLuint program = getProgramFromShaderSource(m_pVertexShaderSource, m_pFragmentShaderSource);
    glUseProgram(program);
    initMesh(program);

    GLuint y_textureId;
    GLuint uv_textureId;
    {
        glGenTextures(1, &y_textureId);
        glGenTextures(1, &uv_textureId);

        GLint y_texture_loc = glGetUniformLocation(program, "yTexture");
        GLint uv_texture_loc = glGetUniformLocation(program, "uvTexture");
        printf("TSDEBUG == func %s, line %d y_texture_loc %d, uv_texture_loc %d\n", __func__, __LINE__, y_texture_loc, uv_texture_loc);
        glUniform1i(y_texture_loc, 0);
        glUniform1i(uv_texture_loc, 1);
    }
    GLuint lut_imageTextureId;
    {
        glGenTextures(1, &lut_imageTextureId);
        glBindTexture(GL_TEXTURE_3D, lut_imageTextureId);
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        float *lutData = CUBE;
        printf("TSDEBUG == func %s, line %d lut_imageTextureId %d\n", __func__, __LINE__, lut_imageTextureId);
        glTexImage3D(GL_TEXTURE_3D, 0, GL_RGB16F, NUM_GRID_LUT_SRC, NUM_GRID_LUT_SRC, NUM_GRID_LUT_SRC, 0, GL_RGB, GL_FLOAT, lutData);
        glBindTexture(GL_TEXTURE_3D, 0);

        GLint lut_texture_loc = glGetUniformLocation(program, "lutTexture");
        printf("TSDEBUG == func %s, line %d lut_texture_loc %d\n", __func__, __LINE__, lut_texture_loc);
        glUniform1i(lut_texture_loc, 2);
    }
    GLuint filterRatioHandle = glGetUniformLocation(program, "filterRatio");
    printf("TSDEBUG == func %s, line %d filterRatioHandle %d\n", __func__, __LINE__, filterRatioHandle);
    {
        //uvoffset
        int bitLength = IMAGE_BIT_LENGTH;
        int uBounds = 1 << bitLength;
        const uint16_t uvoffset = uBounds >> 1;
        GLuint uvofsetHandle = glGetUniformLocation(program, "uvoffset");
        printf("TSDEBUG == func %s, line %d uvofset %d\n", __func__, __LINE__, uvofsetHandle);
        glUniform1ui(uvofsetHandle, uvoffset);
    }
    int w, h;
    uint16_t *srcYUV = loadImage(&w, &h);  /* When you implement driver, the srcYUV should be picked from IPEOUT. */
    uint16_t *srcUVaddr = srcYUV + w * h;
    uint8_t* dstYUV = new uint8_t[windowWidth * windowHeight * 3];

    GLuint FBO;
    GLuint colorAttachmentId;
    {
        glGenFramebuffers(1, &FBO);
        glBindFramebuffer(GL_FRAMEBUFFER, FBO);
        glGenTextures(1, &colorAttachmentId);
        glBindTexture(GL_TEXTURE_2D, colorAttachmentId);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, windowWidth, windowHeight, 0, GL_RGB, GL_UNSIGNED_BYTE, 0);
        glBindTexture(GL_TEXTURE_2D, 0);

        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, colorAttachmentId, 0);
        uint32_t fboStatus = glCheckFramebufferStatus(GL_FRAMEBUFFER);
        if(fboStatus != GL_FRAMEBUFFER_COMPLETE){
            printf("func %s, line %d, OpenGL renderFramebuffer NOT COMPLETE: %d", __func__, __LINE__, fboStatus);
        }
        glBindFramebuffer(GL_FRAMEBUFFER, 0);
    }

    while (!glfwWindowShouldClose(pWindow)) {
        glfwPollEvents();
        glBindFramebuffer(GL_FRAMEBUFFER, 0);
        glViewport(0, 0, windowWidth, windowHeight);
        glScissor(0, 0, windowWidth, windowHeight);
        glClearColor(1.f, 0.f, 0.f, 1.f);
        glClear(GL_COLOR_BUFFER_BIT);

        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, y_textureId);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_R16UI, w, h, 0, GL_RED_INTEGER, GL_UNSIGNED_SHORT, srcYUV);

        glActiveTexture(GL_TEXTURE1);
        glBindTexture(GL_TEXTURE_2D, uv_textureId);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RG16UI, w / 2, h / 2, 0, GL_RG_INTEGER, GL_UNSIGNED_SHORT, srcUVaddr);

        glActiveTexture(GL_TEXTURE2);
        glBindTexture(GL_TEXTURE_3D, lut_imageTextureId);

        glUniform1f(filterRatioHandle, filterRatio);

        glBindVertexArray(VAO);
        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
        glBindVertexArray(0);
        glReadPixels(0, 0, windowWidth, windowHeight, GL_RGB, GL_UNSIGNED_BYTE, dstYUV);

        //glBindFramebuffer(GL_FRAMEBUFFER, 0);
        glfwSwapBuffers(pWindow);
    }
    for (int i = 0; i < 4; i++) {
        printf("{%d,%d,%d}\n", dstYUV[i * 3], dstYUV[i * 3 + 1], dstYUV[i * 3 + 2]);
    }

    writeRgb("output.rgb", windowWidth, windowHeight, windowWidth * windowHeight * 3 * sizeof(uint8_t), dstYUV, true);

    glDeleteTextures(1, &y_textureId);
    glDeleteTextures(1, &uv_textureId);
    glDeleteTextures(1, &lut_imageTextureId);
    glDeleteTextures(1, &colorAttachmentId);
    glDeleteFramebuffers(1, &FBO);
    glDeleteProgram(program);
    delete[] dstYUV;
    delete[] srcYUV;
    glfwDestroyWindow(pWindow);
    glfwTerminate();
    return 0;
}